package com.sali.utils.limiter;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

/**
 * @description:
 * @author: zhang hang
 * @date: 2022/7/19 14:20.
 * @version: 1.0
 */
public class CustomTokenAlgorithm {

    private ArrayBlockingQueue<String> blockingQueue;

    /**
     * 容量大小
     */
    private int limit;

    /**
     * 令牌的产生间隙
     */
    private int period;

    /**
     * 令牌每次产生的个数
     */
    private int amount;

    /**
     * 令牌桶算法初始化
     */
    public CustomTokenAlgorithm(int limit, int period, int amount) {
        this.limit = limit;
        this.period = period;
        this.amount = amount;
        blockingQueue = new ArrayBlockingQueue<>(limit);
    }

    /**
     * 令牌桶初始化
     */
    public void init() {
        for (int i = 0; i < limit; i++) {
            blockingQueue.offer("port");
        }
    }

    /**
     * 令牌桶添加令牌
     */
    private void addToken(int amount) {
        for (int i = 0; i < amount; i++) {
            blockingQueue.offer("port");
        }
    }

    /**
     * 获取令牌
     */
    public boolean tryAcquire() {
        // 队首元素出队
        return blockingQueue.poll() != null;
    }

    /**
     * 令牌桶生产令牌(线程池)
     */
    public void start(Object lock) {

        new ScheduledThreadPoolExecutor(1, new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread thread = new Thread(r);
                return thread;
            }
        }).scheduleAtFixedRate(() -> {
            synchronized (lock) {
                addToken(this.amount);
                lock.notify();
            }
        }, 500, this.period, TimeUnit.MILLISECONDS);

    }


}
